#!/usr/bin/env python3
import rospy
import numpy as np
from grid_map_msgs.msg import GridMap
from std_msgs.msg import Float32MultiArray  # 替换为实际使用的类型
from geometry_msgs.msg import Pose, Point
import matplotlib.pyplot as plt

class ElevationMapSubscriber:
    def __init__(self):
        rospy.init_node('elevation_map_processor', anonymous=True)
        
        # 订阅裁剪后的高程图
        rospy.Subscriber("/elevation_map_clip", GridMap, self.map_callback)
        self.elevation_data = None
        self.map_info = {}
        rospy.loginfo("Elevation Map Processor started. Waiting for data...")

    def map_callback(self, msg):
        # 打印可用层名称用于调试
        rospy.loginfo(f"Available layers: {msg.layers}")
        
        if 'elevation' not in msg.layers:
            rospy.logwarn("No 'elevation' layer found in GridMap!")
            return
        
        layer_idx = msg.layers.index('elevation')
        array_data = msg.data[layer_idx]

        # 获取地图尺寸（关键修正）
        if isinstance(array_data, Float32MultiArray):
            # 从 layout.dim 获取维度
            height = array_data.layout.dim[0].size
            width = array_data.layout.dim[1].size
        else:
            rospy.logerr("Invalid data type for elevation layer.")
            return

        # 更新地图元数据
        self.map_info.update({
            'resolution': msg.info.resolution,
            'width': width,
            'height': height,
            'position': {
                'x': msg.info.pose.position.x,
                'y': msg.info.pose.position.y
            }
        })
        rospy.loginfo(f"Received new elevation map: {width}x{height} @ {msg.info.resolution}m/px")
        self.process_elevation_data()

    def process_elevation_data(self):
        """处理高程数据示例"""
        if self.elevation_data is None:
            return
            
        # 1. 访问高程数据示例
        center_x = self.map_info['width'] // 2
        center_y = self.map_info['height'] // 2
        rospy.loginfo(f"Center elevation: {self.elevation_data[center_y, center_x]} meters")
        
        # 2. 计算统计信息
        valid_data = self.elevation_data[~np.isnan(self.elevation_data)]
        if valid_data.size > 0:
            rospy.loginfo(f"Elevation stats: Min={np.min(valid_data):.2f}m, "
                         f"Max={np.max(valid_data):.2f}m, "
                         f"Mean={np.mean(valid_data):.2f}m")
        
        # 3. 可视化 (可选)
        self.visualize_map()

    def visualize_map(self):
        """使用matplotlib可视化高程图"""
        plt.figure(figsize=(10, 8))
        
        # 创建坐标网格
        res = self.map_info['resolution']
        offset_x = self.map_info['position']['x'] - (self.map_info['width'] * res) / 2
        offset_y = self.map_info['position']['y'] - (self.map_info['height'] * res) / 2
        x = np.arange(0, self.map_info['width']) * res + offset_x
        y = np.arange(0, self.map_info['height']) * res + offset_y
        X, Y = np.meshgrid(x, y)
        
        # 绘制高程图
        plt.contourf(X, Y, self.elevation_data, 20, cmap='terrain')
        plt.colorbar(label='Elevation (m)')
        plt.title('Elevation Map')
        plt.xlabel('X position (m)')
        plt.ylabel('Y position (m)')
        plt.tight_layout()
        
        # 保存或显示
        plt.savefig('/tmp/elevation_map.png')
        rospy.loginfo("Saved elevation visualization to /tmp/elevation_map.png")
        plt.show()  # 在GUI环境中取消注释以显示

if __name__ == '__main__':
    try:
        processor = ElevationMapSubscriber()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass